package com.hdu.dwh.controller;

import java.io.OutputStreamWriter;
import java.nio.charset.StandardCharsets;
import java.sql.ResultSet;
import java.util.List;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;

import com.hdu.dwh.pojos.QueryResult;
import com.hdu.dwh.pojos.source_access_pojos.DataWarehouseDatabase;
import com.hdu.dwh.pojos.source_access_pojos.DataWarehouseTable;
import com.hdu.dwh.results.ResponseResult;
import com.hdu.dwh.service.DataWarehouseDatabaseService;
import com.hdu.dwh.service.DataWarehouseTableService;
import com.hdu.dwh.service.TrinoQueryService;

import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import lombok.extern.slf4j.Slf4j;

@RestController
@Slf4j
@Tag(name = "Trino查询", description = "Trino查询接口")
@RequestMapping("/trino_query")
public class TrinoQueryController {
    @Autowired
    private TrinoQueryService trinoQueryService;

    @Autowired
    private DataWarehouseTableService dwhTableService;

    @Autowired
    private DataWarehouseDatabaseService dwhDatabaseService;

    @Operation(summary = "执行TrinoSQL查询")
    @GetMapping("fetchTrinoQueryResult")
    public ResponseResult<QueryResult> fetchTrinoQueryResult(@RequestParam String SQL) throws Exception {
        return ResponseResult.success(trinoQueryService.executeQuery(SQL));
    }

    @Operation(summary = "执行TrinoSQL查询，指定数据库和表")
    @GetMapping("fetchTrinoQueryResultWithDatabaseAndTable/{page}/{pageSize}")
    public ResponseResult<QueryResult> fetchTrinoQueryResultWithDatabaseAndTable(@RequestParam String database,
                                                                                 @RequestParam String table, @PathVariable Integer page, @PathVariable Integer pageSize) throws Exception {
        // 检查分页参数
        if (page == null || page <= 0) {
            page = 1;
        }
        if (pageSize == null || pageSize <= 0 || pageSize > 1000) {
            pageSize = 1000;
        }

        // 计算offset
        int offset = (page - 1) * pageSize;

        QueryResult result = trinoQueryService.queryWithDatabaseAndTable(database, table, pageSize, offset);
        return ResponseResult.success(result);
    }

    @Operation(summary = "导出数据库表数据为CSV文件")
    @GetMapping("export")
    public ResponseEntity<StreamingResponseBody> exportTableToCSV(@RequestParam Long tableId) throws Exception {
        DataWarehouseTable table = dwhTableService.getById(tableId);
        DataWarehouseDatabase database = dwhDatabaseService.getById(table.getDatabaseId());
        int pageSize = 5000;

        String filename = "table_" + table.getTableName() + ".csv";
        StreamingResponseBody stream = outputStream -> {
            OutputStreamWriter writer = new OutputStreamWriter(outputStream, StandardCharsets.UTF_8);
            // 写CSV表头
            try {
                QueryResult headerResult = trinoQueryService.queryWithDatabaseAndTable(
                        database.getDatabaseName(), table.getTableName(), 1, 0);
                writer.write(String.join(",", headerResult.getColumnNames()) + "\n");
            } catch (Exception e) {
                // TODO Auto-generated catch block
                e.printStackTrace();
                throw new RuntimeException(e);
            }

            writer.flush();

            int page = 0;
            boolean hasMore = true;
            while (hasMore) {
                // 伪代码：分页查询
                QueryResult pageResult;
                try {
                    pageResult = this.trinoQueryService.queryWithDatabaseAndTable(
                            database.getDatabaseName(), table.getTableName(), pageSize, page * pageSize);
                } catch (Exception e) {
                    // TODO Auto-generated catch block
                    e.printStackTrace();
                    throw new RuntimeException(e);
                }
                for (List<String> row : pageResult.getRows()) {
                    // 你需要将 row 转为字符串数组
                    writer.write(String.join(",", row) + "\n");
                }
                writer.flush();
                hasMore = pageResult.getRows().size() == pageSize;
                page++;
            }
            writer.close();
        };

        return ResponseEntity.ok()
                .header(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=\"" + filename + "\"")
                .contentType(MediaType.parseMediaType("text/csv"))
                .body(stream);
    }
}
